/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef FUSION_ENGINE_OPTIMIZER_GRAPH_OPTIMIZER_OP_SETTER_OP_SETTER_H_
#define FUSION_ENGINE_OPTIMIZER_GRAPH_OPTIMIZER_OP_SETTER_OP_SETTER_H_

#include "ops_store/op_kernel_info.h"
#include "common/op_slice_info.h"

namespace fe {
using FillupSliceInfo = std::function<Status(ge::OpDescPtr, OpCalcInfo&)>;
using FillupSliceInfoPtr = std::shared_ptr<FillupSliceInfo>;

class OpSetter {
 public:
  explicit OpSetter(const std::string engine_name);
  virtual ~OpSetter();
  Status SetOpInfo(ge::ComputeGraph& graph);
  Status SetOpInfoByNode(ge::NodePtr node_ptr);

  /*
   *  @ingroup fe
   *  @brief   Fill in the corresponding slice information for each node based on the 'SlicePattern' in the ops
   *           info config.
   *  @param   [in|out] node
   *  @return  SUCCESS or FAILED
   */
  Status SetOpSlicePatternInfoByNode(ge::OpDescPtr op_desc_ptr, const OpKernelInfoPtr& op_kernel_info_ptr);

  /*
   *  @ingroup fe
   *  @brief   Set the slice information for the node whose SlicePattern is Elemwise in the ops info config.
   *  @param   [in|out] node | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status FillupElemwiseSliceInfo(ge::OpDescPtr op_desc_ptr, OpCalcInfo &op_calc_info);

  /*
   *  @ingroup fe
   *  @brief   Set the slice information for the node whose SlicePattern is ElemwiseBroadcast in the ops info config.
   *  @param   [in|out] node | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status FillupElemwiseBroadcastSliceInfo(ge::OpDescPtr op_desc_ptr, OpCalcInfo &op_calc_info);

  /*
   *  @ingroup fe
   *  @brief   Set the slice information for the node whose SlicePattern is Broadcast in the ops info config.
   *  @param   [in|out] node | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status FillupBroadcastSliceInfo(ge::OpDescPtr op_desc_ptr, OpCalcInfo &op_calc_info);

  /*
   *  @ingroup fe
   *  @brief   Set the slice information for the node whose SlicePattern is SlidingWindow in the ops info config.
   *  @param   [in|out] node | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status FillupSlidingWindowSliceInfo(ge::OpDescPtr op_desc_ptr, OpCalcInfo &op_calc_info);

  /*
   *  @ingroup fe
   *  @brief   Set the slice information for the node whose SlicePattern is SlidingWindowDeconv in the ops info config.
   *  @param   [in|out] node | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status FillupSlidingWindowDeconvSliceInfo(ge::OpDescPtr op_desc_ptr, OpCalcInfo &op_calc_info);

  /*
   *  @ingroup fe
   *  @brief   Set the slice information for the node whose SlicePattern is CubeMatmul in the ops info config.
   *  @param   [in|out] node | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status FillupCubeMatmulSliceInfo(ge::OpDescPtr op_desc_ptr, OpCalcInfo &op_calc_info);

  /*
   *  @ingroup fe
   *  @brief   Set the slice information for the node whose SlicePattern is Reduce in the ops info config.
   *  @param   [in|out] node | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status FillupReduceSliceInfo(ge::OpDescPtr op_desc_ptr, OpCalcInfo &op_calc_info);

  /*
   *  @ingroup fe
   *  @brief   Set the slice information for the node whose SlicePattern is Resize in the ops info config.
   *  @param   [in|out] node | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status FillupResizeSliceInfo(ge::OpDescPtr op_desc_ptr, OpCalcInfo &op_calc_info);

  /*
   *  @ingroup fe
   *  @brief   Set the slice information for the node whose SlicePattern is Scatter in the ops info config.
   *  @param   [in|out] node | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status FillupScatterSliceInfo(ge::OpDescPtr op_desc_ptr, OpCalcInfo &op_calc_info);

  /*
   *  @ingroup fe
   *  @brief   Set the slice information for the node whose SlicePattern is Segment in the ops info config.
   *  @param   [in|out] node | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status FillupSegmentSliceInfo(ge::OpDescPtr op_desc_ptr, OpCalcInfo &op_calc_info);


 private:
  Status SetOpDescAttr(const OpKernelInfoPtr& op_kernel_info_ptr, const std::string& attr_name,
                       const std::string& ge_attr_name, ge::OpDescPtr op_desc_ptr);
  static Status CheckElemwiseInputAndOutputNum(ge::OpDescPtr op_desc_ptr, bool &has_scalar, size_t &dim_size,
                                               ge::Format &op_output_format);
  /*
   *  @ingroup fe
   *  @brief   Set the slice information corresponding to the elemwise based on the node information.
   *  @param   [in|out] node & has scalar or not | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status SetElemWiseSliceInfo(ge::OpDescPtr op_desc_ptr,
                                     std::vector<AxisSplitMap>& axis_split_maps, bool has_scalar);
  /*
   *  @ingroup fe
   *  @brief   Set the slice information corresponding to the slidingwindow based on the node information.
   *  @param   [in|out] node & has scalar or not | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status SetSlidingWindowSliceInfo(ge::OpDescPtr op_desc_ptr, std::vector<AxisSplitMap>& axis_split_maps);
  /*
   *  @ingroup fe
   *  @brief   Set the slice information corresponding to the slidingwindowdeconv based on the node information.
   *  @param   [in|out] node & has scalar or not | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status SetSlidingWindowDeconvSliceInfo(ge::OpDescPtr op_desc_ptr, std::vector<AxisSplitMap>& axis_split_maps);
  /*
   *  @ingroup fe
   *  @brief   Set the slice information corresponding to the cube matmul based on the node information.
   *  @param   [in|out] node & has scalar or not | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status SetCubeMatmulSliceInfo(ge::OpDescPtr op_desc_ptr, std::vector<AxisSplitMap>& axis_split_maps);
  /*
   *  @ingroup fe
   *  @brief   Set the slice information corresponding to the reduce based on the node information.
   *  @param   [in|out] node & has scalar or not | slice info
   *  @return  SUCCESS or FAILED
   */
  static Status SetReduceSliceInfo(ge::OpDescPtr op_desc_ptr, std::vector<AxisSplitMap>& axis_split_maps);

  static void SetInputSplitInfo(AxisSplitMap& axis_split_map, const int8_t& input_index, const int8_t& input_axis);
  static void SetInputReduceInfo(AxisReduceMap& axis_split_map, const int8_t& input_index, const int8_t& input_axis);
  static void SetOutputSplitInfo(AxisSplitMap& axis_split_map, const int8_t& output_index, const int8_t& output_axis);
  static void SetOutputReduceInfo(AxisReduceMap& axis_reduce_map, const int8_t& output_index,
                                  const OpReduceType& reduce_type, const bool& is_atomic);
  /*
   *  @ingroup fe
   *  @brief   Set the slice information based on the input and output indexes/axes.
   *  @param   [in|out] input and output indexes/axes | slice info
   *  @return  SUCCESS or FAILED
   */
  static void SetMultiAxisSplitMap(AxisSplitMap& axis_split_map, const int8_t& first_index, const int8_t& first_axis,
                                   const int8_t& output_index, const int8_t& output_axis,
                                   const int8_t second_index = -1, const int8_t second_axis = -1);
  std::string engine_name_;

  static const std::map<SlicePattern, FillupSliceInfoPtr> split_info_map_;
  static const std::map<string, string> attr_map_;
};
}  // namespace fe
#endif  // FUSION_ENGINE_OPTIMIZER_GRAPH_OPTIMIZER_OP_SETTER_OP_SETTER_H_